# 导入 OpenCV 库，用于图像处理，如读取、修改和保存图像等操作
import cv2
# 导入 NumPy 库，用于数值计算，在图像处理中可高效处理数组
import numpy as np
# 导入 PyTorch 深度学习框架，用于构建和训练神经网络
import torch
# 从 PyTorch 中导入 Dataset 类，用于自定义数据集
from torch.utils.data import Dataset
# 从 PyTorch 中导入 transforms 模块，可用于图像变换，但当前未使用
from torchvision import transforms
# 从 PyTorch 中导入 ImageFolder 类，用于从文件夹中加载图像数据集
from torchvision.datasets import ImageFolder


def image_square_pad(image, color=(0, 0, 0)):
    """
    Pad image to square

    Args:
        image (numpy.ndarray): 输入的图像数组
        color (tuple, optional): 填充颜色，默认为黑色 (0, 0, 0)

    Returns:
        tuple: 填充后的图像数组和填充信息 [pad_top, pad_bottom, pad_left, pad_right]
    """
    # 获取图像的高度、宽度和通道数
    im_h, im_w, _ = image.shape
    # 找出图像的最大边长
    max_edge = max(im_h, im_w)
    # 计算顶部填充量
    pad_top = (max_edge - im_h) // 2
    # 计算底部填充量
    pad_bottom = max_edge - im_h - pad_top
    # 计算左侧填充量
    pad_left = (max_edge - im_w) // 2
    # 计算右侧填充量
    pad_right = max_edge - im_w - pad_left
    
    # 使用 OpenCV 的 copyMakeBorder 函数对图像进行填充
    image = cv2.copyMakeBorder(image, pad_top, pad_bottom, pad_left, pad_right, cv2.BORDER_CONSTANT, value=color)
    return image, [pad_top, pad_bottom, pad_left, pad_right]


def rotate_image(image, angle):
    # 获取图片的高度和宽度
    (h, w) = image.shape[:2]

    # 计算图像的中心
    center = (w // 2, h // 2)

    # 获取旋转矩阵
    M = cv2.getRotationMatrix2D(center, angle, 1.0)

    # 计算旋转后图像的大小
    abs_cos = abs(M[0, 0])
    abs_sin = abs(M[0, 1])
    new_w = int(h * abs_sin + w * abs_cos)
    new_h = int(h * abs_cos + w * abs_sin)

    # 让旋转后的图像能够完全显示
    M[0, 2] += (new_w / 2) - center[0]
    M[1, 2] += (new_h / 2) - center[1]

    # 旋转并填充多出来的区域为 0
    rotated_image = cv2.warpAffine(image, M, (new_w, new_h), borderMode=cv2.BORDER_CONSTANT, borderValue=(0, 0, 0))

    return rotated_image


class SwitchDataset(Dataset):
    def __init__(self, root, size=224):
        """
        初始化 SwitchDataset 类

        Args:
            root (str): 数据集的根目录
            size (int, optional): 图像的目标大小，默认为 224
        """
        # 保存数据集的根目录
        self.root = root
        # 使用 ImageFolder 加载数据集
        self.data = ImageFolder(root)
        
        # 保存图像的目标大小
        self.size = size
        # 获取数据集中的类别名称列表
        self.class_names = self.data.classes
        # 获取类别到索引的映射字典
        self.class_to_idx = self.data.class_to_idx
        # 计算数据集中的类别数量
        self.num_classes = len(self.class_names)

    def __len__(self):
        """
        返回数据集的长度

        Returns:
            int: 数据集中样本的数量
        """
        return len(self.data)
    
    def __getitem__(self, idx):
        """
        根据索引获取数据集中的样本

        Args:
            idx (int): 样本的索引

        Returns:
            tuple: 处理后的图像张量和对应的标签张量
        """
        # 根据索引获取图像和标签
        image, label = self.data[idx]
        # 将图像转换为 NumPy 数组
        image = np.array(image)
        
        # 以 0.5 的概率对图像进行高斯模糊处理，且图像的最小边长需大于 50
        if np.random.rand() < 0.5 and min(image.shape[:2]) > 50:
            image = cv2.GaussianBlur(image, (3, 3), sigmaX=0.5)
        
        # 以 0.5 的概率对图像进行水平翻转
        if np.random.rand() < 0.5:
            image = cv2.flip(image, 1)
        
        # 将图像填充为正方形
        image, _ = image_square_pad(image)
        
        # 随机生成一个 0 到 4 之间的角度
        angle = np.random.randint(0, 5)
        # 对图像进行旋转
        image = rotate_image(image, angle)

        # 将图像调整为指定大小
        image = cv2.resize(image, (self.size, self.size))

        # 对图像进行归一化处理
        image = image / 255.0 * 2 - 1

        # 将图像数组转换为 PyTorch 张量，并调整维度顺序
        image = torch.from_numpy(image).permute(2, 0, 1).float()
        # 将标签转换为 PyTorch 张量
        label = torch.tensor(label).long()
        return image, label

    def get_class_names(self):
        """
        获取数据集中的类别名称列表

        Returns:
            list: 类别名称列表
        """
        return self.class_names
    
    def get_class_to_idx(self):
        """
        获取类别到索引的映射字典

        Returns:
            dict: 类别到索引的映射字典
        """
        return self.class_to_idx